import cv2
import numpy as np


# 棋盘格图的行数和列数
cbraw = 6
cbcol = 7

# 上一节中测得的相机内参数矩阵
camera_matrix = np.array((
    [1.40787017e+03, 0.00000000e+00, 9.48648860e+02],
    [0.00000000e+00, 1.45577464e+03, 5.63655869e+02],
    [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]), dtype=np.double)


def get_corners_from_img(path):
    """
    从一张图片中提取角点
    :param path: 图片路径
    :return: img, ret, corners， img是读取到的图片，ret是结果，corners是角点信息
    """
    img = cv2.imread(path)  # source image

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 转灰度

    # 寻找角点，存入corners，ret是找到角点的flag
    ret, corners = cv2.findChessboardCorners(gray, (cbraw, cbcol), None)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)

    return img, ret, corners2


def draw_corners_id(img, ret, corners2, window_name):
    """
    在棋盘格图片上标记获得的角点的id，以查看角点是否是对应的上的
    :param img: 在哪张图上画
    :param ret: 角点检测结果
    :param corners2: 角点信息
    :param window_name: 窗口名称
    :return: 
    """
    count = 0
    for corner in corners2:
        a, b = corner[0]
        a = int(a)
        b = int(b)

        cv2.putText(img, '(%s)' % (count), (a, b), cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 0), 1)
        print('id={}:({},{})'.format(count, a, b))
        count += 1

        img = cv2.drawChessboardCorners(img, (cbraw, cbcol), corners2, ret)
        cv2.imshow(window_name, img)
        # cv2.imwrite('{}.jpg'.format(window_name), img)


def run():
    # 1.获取左右两幅图像的特征点。
    img_l, ret_l, corners_l = get_corners_from_img('./imgs3/l.jpg')
    img_r, ret_r, corners_r = get_corners_from_img('./imgs3/r.jpg')
    # 查看检测到的角点是否对应
    # draw_corners_id(img_l, ret_l, corners_l, 'left')
    # draw_corners_id(img_r, ret_r, corners_r, 'right')
    # print(ret_l, corners_l, 'left')
    # print(ret_r, corners_r, 'right')

    # 2. 使用findEssentialMat计算本质矩阵E
    E, mask = cv2.findEssentialMat(corners_l, corners_r,
                                   cameraMatrix=camera_matrix, method=cv2.RANSAC,
                                   threshold=1, prob=0.999
                                   )
    print('E: ', E)
    # 3. 使用recoverPose 从E中分解出两个相机相对位姿
    retval2, R, t, mask = cv2.recoverPose(E, corners_l, corners_r, camera_matrix)
    print("retval2:", retval2)
    print("R:", R)
    print("t:", t)
    dstR, _ = cv2.Rodrigues(R)
    print("dstR:", dstR)
    pass


if __name__ == '__main__':
    run()

    cv2.waitKey(50000)
